{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "47e7e884",
   "metadata": {},
   "outputs": [],
   "source": [
    "deployment = \"gpt4\"\n",
    "model = \"gpt-4\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "311fcae6",
   "metadata": {},
   "source": [
    "# 大模型辅助设计"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "66522ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "def work_on(input):\n",
    "    response = openai.ChatCompletion.create(\n",
    "        engine=deployment, # 如果直接访问OpenAI GPT服务的同学，这里不要使用engine这个参数，要使用model，如： model=“gpt-4”\n",
    "        model=model,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a senior software engineer.\"},   \n",
    "            {\"role\": \"user\", \"content\": input}\n",
    "        ],\n",
    "        temperature = 0.9, \n",
    "        max_tokens = 500\n",
    "      )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f49fb28a",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "Design a microservice according to the following requirements:\n",
    "The service will process the user's input and store the result. \n",
    "After processing a token will be returned to the user, it would be used to get the processing result.\n",
    "Consider on the following design principles:\n",
    "1. concerns separation, try you best to separate the functions into the small modules.\n",
    "2. single responsibility \n",
    "3. interface oriented\n",
    "make sure the extensibility and reusability of the design, represent your design (with main abstractions) with PlantUML script\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "87b6ead7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Below is a basic design of the microservice:\n",
      "\n",
      "The microservice will be split into the following modules:\n",
      "\n",
      "1. **InputProcessor:** This module will be responsible for processing the user's input. Interface-oriented design will be followed to ensure extensibility and reusability.\n",
      "\n",
      "2. **TokenGenerator:** This module will generate a unique token for each processed input. This design allows for separation of concerns and follows the single responsibility principle.\n",
      "\n",
      "3. **ResultStore:** This module will store the processed result against the generated token. This design also follows the single responsibility principle.\n",
      "\n",
      "4. **TokenValidator:** This module will validate the provided token and fetch the corresponding result.\n",
      "\n",
      "Here is the PlantUML script representing the above design:\n",
      "\n",
      "```plantuml\n",
      "@startuml\n",
      "title Microservice Design\n",
      "\n",
      "class InputProcessor {\n",
      "  +process(input: Input): Result\n",
      "}\n",
      "\n",
      "class TokenGenerator {\n",
      "  +generateToken(result: Result): Token\n",
      "}\n",
      "\n",
      "class ResultStore {\n",
      "  +store(token: Token, result: Result)\n",
      "  +getResult(token: Token): Result\n",
      "}\n",
      "\n",
      "class TokenValidator {\n",
      "  +validate(token: Token): Boolean\n",
      "}\n",
      "\n",
      "InputProcessor->TokenGenerator: result\n",
      "TokenGenerator->ResultStore: token\n",
      "TokenValidator->ResultStore: validated token\n",
      "@enduml\n",
      "```\n",
      "\n",
      "In this design, every module has a single responsibility which makes the design extensible and reusable. The communication between the modules is done through interfaces thus ensuring that the implementation details of one module do not affect the others. This promotes concerns separation. The user's input is processed and a token is generated which is used to retrieve the result later. All this functionality is encapsulated within separate modules.\n"
     ]
    }
   ],
   "source": [
    "print(work_on(prompt))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a18d6c98",
   "metadata": {},
   "source": [
    "# 生成Plugin代码及Skeleton单元测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3b473a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import re\n",
    "\n",
    "def find_declarations(directory, json_string):\n",
    "    # Parse the JSON string\n",
    "    json_data = json.loads(json_string)\n",
    "    package_name = json_data[\"package_name\"]\n",
    "    element_name = json_data[\"element\"]\n",
    "\n",
    "    # Regular expression to match structure/interface declarations\n",
    "    regex = re.compile(r\"(type\\s+\" + element_name + r\"\\s+(struct|interface)\\s+{\\s+[^}]*\\s+})\")\n",
    "\n",
    "    # Read the source files in the directory\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if file.endswith(\".go\"):\n",
    "                with open(os.path.join(root, file), \"r\") as f:\n",
    "                    content = f.read()\n",
    "                    # Check if the file is in the correct package\n",
    "                    if \"package \" + package_name in content:\n",
    "                        # Search for the structure/interface declaration\n",
    "                        match = regex.search(content)\n",
    "                        if match:\n",
    "                            return match.group(1)\n",
    "\n",
    "    return \"Not Found\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c303715c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_file(file_path):\n",
    "    with open(file_path, 'r') as file:\n",
    "        content = file.read()\n",
    "    return content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e709df39",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import initialize_agent, Tool\n",
    "from langchain.chat_models import AzureChatOpenAI\n",
    "from langchain.chains import LLMRequestsChain\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.llms import AzureOpenAI\n",
    "import time\n",
    "import os\n",
    "import json\n",
    "import re\n",
    "\n",
    "src_dir = \"./llm_friendly_code\"\n",
    "\n",
    "def find_declarations(directory, json_string):\n",
    "    # Parse the JSON string\n",
    "    json_data = json.loads(json_string)\n",
    "    package_name = json_data[\"package_name\"]\n",
    "    element_name = json_data[\"element\"]\n",
    "\n",
    "    # Regular expression to match structure/interface declarations\n",
    "    regex = re.compile(r\"(type\\s+\" + element_name + r\"\\s+(struct|interface)\\s+{\\s+[^}]*\\s+})\")\n",
    "\n",
    "    # Read the source files in the directory\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if file.endswith(\".go\"):\n",
    "                with open(os.path.join(root, file), \"r\") as f:\n",
    "                    content = f.read()\n",
    "                    # Check if the file is in the correct package\n",
    "                    if \"package \" + package_name in content:\n",
    "                        # Search for the structure/interface declaration\n",
    "                        match = regex.search(content)\n",
    "                        if match:\n",
    "                            return match.group(1)\n",
    "\n",
    "    return \"Not Found\"\n",
    "\n",
    "def get_dep(json_string):\n",
    "    return find_declarations(src_dir, json_string)\n",
    "\n",
    "llm = AzureChatOpenAI(deployment_name=deployment, temperature=0)\n",
    "\n",
    "# 将“依赖获取”工具添加到 ReAct 过程中    \n",
    "tools = [\n",
    "    Tool(\n",
    "        name=\"find_declarations\", func=get_dep,\n",
    "        description=\"\"\"Only invoke it when you need to find the declaration of the interface/struct/method that is not in the given code. \n",
    "        The input is a JSON string includes package name and the struct/interface name, \n",
    "          and its the format {{\"package_name\":<package name>,\"element\":<struct/interface>}}\n",
    "        \"\"\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "96444d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_code(input):\n",
    "    agent = initialize_agent(tools, llm, agent=\"zero-shot-react-description\", handle_parsing_errors=True, max_iterations=10, verbose=True)\n",
    "    print(agent.run(input))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6bde3501",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mFirst, I need to know what methods are in the demo.TokenCreator interface. \n",
      "Action: find_declarations\n",
      "Action Input: {\"package_name\":\"demo\",\"element\":\"TokenCreator\"}\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mtype TokenCreator interface {\n",
      "\tCreateToken(data string) string\n",
      "}\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mI now know the final answer\n",
      "Final Answer: \n",
      "Here is the implementation of the demo.TokenCreator interface in Go:\n",
      "\n",
      "```go\n",
      "package demo\n",
      "\n",
      "import (\n",
      "\t\"github.com/google/uuid\"\n",
      ")\n",
      "\n",
      "type MyTokenCreator struct{}\n",
      "\n",
      "func (m *MyTokenCreator) CreateToken(data string) string {\n",
      "\treturn uuid.New().String() + data\n",
      "}\n",
      "```\n",
      "\n",
      "In this code, we define a struct `MyTokenCreator` that will implement the `TokenCreator` interface. We then define a method `CreateToken` on `MyTokenCreator` that takes a string as input and returns a string. This method generates a new UUID and appends the input data to it to create the token.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Here is the implementation of the demo.TokenCreator interface in Go:\n",
      "\n",
      "```go\n",
      "package demo\n",
      "\n",
      "import (\n",
      "\t\"github.com/google/uuid\"\n",
      ")\n",
      "\n",
      "type MyTokenCreator struct{}\n",
      "\n",
      "func (m *MyTokenCreator) CreateToken(data string) string {\n",
      "\treturn uuid.New().String() + data\n",
      "}\n",
      "```\n",
      "\n",
      "In this code, we define a struct `MyTokenCreator` that will implement the `TokenCreator` interface. We then define a method `CreateToken` on `MyTokenCreator` that takes a string as input and returns a string. This method generates a new UUID and appends the input data to it to create the token.\n"
     ]
    }
   ],
   "source": [
    "write_code(\n",
    "\"\"\"\n",
    " You are a very senior golang programmer.\n",
    " Write the golang code to implement the demo.TokenCreator interface.\n",
    " The implementation is to generate the token with uuid and its package is demo\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f0e79313",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mFirst, I need to know what methods are in the demo.Storage interface.\n",
      "Action: find_declarations\n",
      "Action Input: {\"package_name\":\"demo\",\"element\":\"Storage\"}\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mtype Storage interface {\n",
      "\t// RetiveData is to retrive the data by the associated token.\n",
      "\tRetieveData(token string) (string, error)\n",
      "\n",
      "\t// StoreData is to persist the data,\n",
      "\t// input paramters:\n",
      "\t//   token is used to retrive the associated data\n",
      "\t//\n",
      "\tStoreData(token string, data string) error\n",
      "}\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mI now know the methods in the demo.Storage interface. I need to implement these methods in a struct that represents local memory storage.\n",
      "Final Answer: \n",
      "Here is the Go code to implement the demo.Storage interface based on local memory:\n",
      "\n",
      "```go\n",
      "package main\n",
      "\n",
      "import (\n",
      "\t\"errors\"\n",
      "\t\"sync\"\n",
      ")\n",
      "\n",
      "type LocalMemoryStorage struct {\n",
      "\tdata map[string]string\n",
      "\tmu   sync.RWMutex\n",
      "}\n",
      "\n",
      "func NewLocalMemoryStorage() *LocalMemoryStorage {\n",
      "\treturn &LocalMemoryStorage{\n",
      "\t\tdata: make(map[string]string),\n",
      "\t}\n",
      "}\n",
      "\n",
      "func (l *LocalMemoryStorage) RetieveData(token string) (string, error) {\n",
      "\tl.mu.RLock()\n",
      "\tdefer l.mu.RUnlock()\n",
      "\n",
      "\tdata, ok := l.data[token]\n",
      "\tif !ok {\n",
      "\t\treturn \"\", errors.New(\"data not found\")\n",
      "\t}\n",
      "\n",
      "\treturn data, nil\n",
      "}\n",
      "\n",
      "func (l *LocalMemoryStorage) StoreData(token string, data string) error {\n",
      "\tl.mu.Lock()\n",
      "\tdefer l.mu.Unlock()\n",
      "\n",
      "\tl.data[token] = data\n",
      "\treturn nil\n",
      "}\n",
      "```\n",
      "\n",
      "In this code, I created a struct `LocalMemoryStorage` that has a map to store data and a RWMutex for concurrent safe operations. The `NewLocalMemoryStorage` function is used to create a new instance of `LocalMemoryStorage`. The `RetieveData` method retrieves data from the map using the provided token. If the data is not found, it returns an error. The `StoreData` method stores data in the map with the provided token.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Here is the Go code to implement the demo.Storage interface based on local memory:\n",
      "\n",
      "```go\n",
      "package main\n",
      "\n",
      "import (\n",
      "\t\"errors\"\n",
      "\t\"sync\"\n",
      ")\n",
      "\n",
      "type LocalMemoryStorage struct {\n",
      "\tdata map[string]string\n",
      "\tmu   sync.RWMutex\n",
      "}\n",
      "\n",
      "func NewLocalMemoryStorage() *LocalMemoryStorage {\n",
      "\treturn &LocalMemoryStorage{\n",
      "\t\tdata: make(map[string]string),\n",
      "\t}\n",
      "}\n",
      "\n",
      "func (l *LocalMemoryStorage) RetieveData(token string) (string, error) {\n",
      "\tl.mu.RLock()\n",
      "\tdefer l.mu.RUnlock()\n",
      "\n",
      "\tdata, ok := l.data[token]\n",
      "\tif !ok {\n",
      "\t\treturn \"\", errors.New(\"data not found\")\n",
      "\t}\n",
      "\n",
      "\treturn data, nil\n",
      "}\n",
      "\n",
      "func (l *LocalMemoryStorage) StoreData(token string, data string) error {\n",
      "\tl.mu.Lock()\n",
      "\tdefer l.mu.Unlock()\n",
      "\n",
      "\tl.data[token] = data\n",
      "\treturn nil\n",
      "}\n",
      "```\n",
      "\n",
      "In this code, I created a struct `LocalMemoryStorage` that has a map to store data and a RWMutex for concurrent safe operations. The `NewLocalMemoryStorage` function is used to create a new instance of `LocalMemoryStorage`. The `RetieveData` method retrieves data from the map using the provided token. If the data is not found, it returns an error. The `StoreData` method stores data in the map with the provided token.\n"
     ]
    }
   ],
   "source": [
    "write_code(\n",
    "\"\"\"\n",
    " You are a very senior golang programmer.\n",
    " Write the golang code to implement the demo.Storage interface.\n",
    " The implementation is based on local memory.\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "74f17d93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mFirst, I need to know what the demo.Processor interface looks like in order to implement it. \n",
      "Action: find_declarations\n",
      "Action Input: {\"package_name\":\"demo\",\"element\":\"Processor\"} \u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mtype Processor interface {\n",
      "\t// Process is to the raw data\n",
      "\tProcess(raw string) (string, error)\n",
      "}\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mNow I know the structure of the demo.Processor interface. It has a single method, Process, which takes a raw string as input and returns a string and an error. I can now write the implementation of this interface. The implementation should reverse the input string.\n",
      "Final Answer: Here is the Go code to implement the demo.Processor interface:\n",
      "\n",
      "```go\n",
      "package main\n",
      "\n",
      "import (\n",
      "\t\"errors\"\n",
      "\t\"strings\"\n",
      ")\n",
      "\n",
      "type ReverseProcessor struct{}\n",
      "\n",
      "func (rp *ReverseProcessor) Process(raw string) (string, error) {\n",
      "\tif raw == \"\" {\n",
      "\t\treturn \"\", errors.New(\"input string is empty\")\n",
      "\t}\n",
      "\n",
      "\trunes := []rune(raw)\n",
      "\tfor i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {\n",
      "\t\trunes[i], runes[j] = runes[j], runes[i]\n",
      "\t}\n",
      "\n",
      "\treturn string(runes), nil\n",
      "}\n",
      "```\n",
      "\n",
      "In this code, I've created a struct called ReverseProcessor which implements the demo.Processor interface. The Process method takes a raw string, checks if it's empty and if it's not, it reverses the string and returns it. If the string is empty, it returns an error.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "Here is the Go code to implement the demo.Processor interface:\n",
      "\n",
      "```go\n",
      "package main\n",
      "\n",
      "import (\n",
      "\t\"errors\"\n",
      "\t\"strings\"\n",
      ")\n",
      "\n",
      "type ReverseProcessor struct{}\n",
      "\n",
      "func (rp *ReverseProcessor) Process(raw string) (string, error) {\n",
      "\tif raw == \"\" {\n",
      "\t\treturn \"\", errors.New(\"input string is empty\")\n",
      "\t}\n",
      "\n",
      "\trunes := []rune(raw)\n",
      "\tfor i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {\n",
      "\t\trunes[i], runes[j] = runes[j], runes[i]\n",
      "\t}\n",
      "\n",
      "\treturn string(runes), nil\n",
      "}\n",
      "```\n",
      "\n",
      "In this code, I've created a struct called ReverseProcessor which implements the demo.Processor interface. The Process method takes a raw string, checks if it's empty and if it's not, it reverses the string and returns it. If the string is empty, it returns an error.\n"
     ]
    }
   ],
   "source": [
    "write_code(\n",
    "\"\"\"\n",
    " You are a very senior golang programmer.\n",
    " Write the golang code to implement the demo.Processor interface.\n",
    " The implementation is to reverse the input content\n",
    "\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "73490b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_unit_tests(filepath):\n",
    "    agent = initialize_agent(tools, llm, agent=\"zero-shot-react-description\", handle_parsing_errors=True, max_iterations=10, verbose=True)\n",
    "    return (agent.run(\n",
    "\n",
    "read_file(filepath) +        \n",
    "\"\"\"\\n\n",
    "    ----\\n\n",
    "你是一个资深软件工程师，请为上面的代码编写单元测试，按下面步骤进行：\n",
    "\n",
    "1. 测试程序与被测试程序在同一 package 中\n",
    "\n",
    "2. 分析需要测试的源码\n",
    "\n",
    "3. 先在遗留代码中查找写测试所需要的接口的声明，便于生成 mock 对象\n",
    "\n",
    "根据以上输出单元测试程序\n",
    "    \n",
    "    \\n ---\"\"\" + read_file(filepath)\n",
    "     ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4a8e5683",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mThe code provided is a service implementation with two methods: Process and Retrieve. These methods use three interfaces: Processor, TokenCreator, and Storage. To write unit tests for this service, I need to create mock implementations of these interfaces. However, the declarations of these interfaces are not provided in the code. I need to find the declarations of these interfaces first.\n",
      "\n",
      "Action: find_declarations\n",
      "Action Input: {\"package_name\":\"demo\",\"element\":\"Processor\"}\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mtype Processor interface {\n",
      "\t// Process is to the raw data\n",
      "\tProcess(raw string) (string, error)\n",
      "}\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mNow I have the declaration of the Processor interface. I need to find the declarations of the TokenCreator and Storage interfaces next.\n",
      "\n",
      "Action: find_declarations\n",
      "Action Input: {\"package_name\":\"demo\",\"element\":\"TokenCreator\"}\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mtype TokenCreator interface {\n",
      "\tCreateToken(data string) string\n",
      "}\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mNow I have the declaration of the TokenCreator interface. I need to find the declaration of the Storage interface next.\n",
      "\n",
      "Action: find_declarations\n",
      "Action Input: {\"package_name\":\"demo\",\"element\":\"Storage\"}\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mtype Storage interface {\n",
      "\t// RetiveData is to retrive the data by the associated token.\n",
      "\tRetieveData(token string) (string, error)\n",
      "\n",
      "\t// StoreData is to persist the data,\n",
      "\t// input paramters:\n",
      "\t//   token is used to retrive the associated data\n",
      "\t//\n",
      "\tStoreData(token string, data string) error\n",
      "}\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mNow I have the declarations of all the interfaces used in the ProcessingServiceImpl. I can now write the unit tests for the Process and Retrieve methods of the ProcessingServiceImpl. \n",
      "\n",
      "Final Answer:\n",
      "\n",
      "```go\n",
      "package demo\n",
      "\n",
      "import (\n",
      "\t\"testing\"\n",
      "\t\"errors\"\n",
      "\t\"github.com/stretchr/testify/mock\"\n",
      "\t\"github.com/stretchr/testify/assert\"\n",
      ")\n",
      "\n",
      "type MockProcessor struct {\n",
      "\tmock.Mock\n",
      "}\n",
      "\n",
      "func (m *MockProcessor) Process(raw string) (string, error) {\n",
      "\targs := m.Called(raw)\n",
      "\treturn args.String(0), args.Error(1)\n",
      "}\n",
      "\n",
      "type MockTokenCreator struct {\n",
      "\tmock.Mock\n",
      "}\n",
      "\n",
      "func (m *MockTokenCreator) CreateToken(data string) string {\n",
      "\targs := m.Called(data)\n",
      "\treturn args.String(0)\n",
      "}\n",
      "\n",
      "type MockStorage struct {\n",
      "\tmock.Mock\n",
      "}\n",
      "\n",
      "func (m *MockStorage) RetieveData(token string) (string, error) {\n",
      "\targs := m.Called(token)\n",
      "\treturn args.String(0), args.Error(1)\n",
      "}\n",
      "\n",
      "func (m *MockStorage) StoreData(token string, data string) error {\n",
      "\targs := m.Called(token, data)\n",
      "\treturn args.Error(0)\n",
      "}\n",
      "\n",
      "func TestProcessingServiceImpl_Process(t *testing.T) {\n",
      "\tmockProcessor := new(MockProcessor)\n",
      "\tmockTokenCreator := new(MockTokenCreator)\n",
      "\tmockStorage := new(MockStorage)\n",
      "\n",
      "\tservice := NewProcessingService(mockProcessor, mockTokenCreator, mockStorage)\n",
      "\n",
      "\traw := \"raw data\"\n",
      "\tprocessed := \"processed data\"\n",
      "\ttoken := \"token\"\n",
      "\n",
      "\tmockProcessor.On(\"Process\", raw).Return(processed, nil)\n",
      "\tmockTokenCreator.On(\"CreateToken\", processed).Return(token)\n",
      "\tmockStorage.On(\"StoreData\", token, processed).Return(nil)\n",
      "\n",
      "\tresult, err := service.Process(raw)\n",
      "\n",
      "\tassert.NoError(t, err)\n",
      "\tassert.Equal(t, token, result)\n",
      "\n",
      "\tmockProcessor.AssertExpectations(t)\n",
      "\tmockTokenCreator.AssertExpectations(t)\n",
      "\tmockStorage.AssertExpectations(t)\n",
      "}\n",
      "\n",
      "func TestProcessingServiceImpl_Retrieve(t *testing.T) {\n",
      "\tmockProcessor := new(MockProcessor)\n",
      "\tmockTokenCreator := new(MockTokenCreator)\n",
      "\tmockStorage := new(MockStorage)\n",
      "\n",
      "\tservice := NewProcessingService(mockProcessor, mockTokenCreator, mockStorage)\n",
      "\n",
      "\ttoken := \"token\"\n",
      "\tdata := \"data\"\n",
      "\n",
      "\tmockStorage.On(\"RetieveData\", token).Return(data, nil)\n",
      "\n",
      "\tresult, err := service.Retrieve(token)\n",
      "\n",
      "\tassert.NoError(t, err)\n",
      "\tassert.Equal(t, data, result)\n",
      "\n",
      "\tmockStorage.AssertExpectations(t)\n",
      "}\n",
      "```\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "```go\n",
      "package demo\n",
      "\n",
      "import (\n",
      "\t\"testing\"\n",
      "\t\"errors\"\n",
      "\t\"github.com/stretchr/testify/mock\"\n",
      "\t\"github.com/stretchr/testify/assert\"\n",
      ")\n",
      "\n",
      "type MockProcessor struct {\n",
      "\tmock.Mock\n",
      "}\n",
      "\n",
      "func (m *MockProcessor) Process(raw string) (string, error) {\n",
      "\targs := m.Called(raw)\n",
      "\treturn args.String(0), args.Error(1)\n",
      "}\n",
      "\n",
      "type MockTokenCreator struct {\n",
      "\tmock.Mock\n",
      "}\n",
      "\n",
      "func (m *MockTokenCreator) CreateToken(data string) string {\n",
      "\targs := m.Called(data)\n",
      "\treturn args.String(0)\n",
      "}\n",
      "\n",
      "type MockStorage struct {\n",
      "\tmock.Mock\n",
      "}\n",
      "\n",
      "func (m *MockStorage) RetieveData(token string) (string, error) {\n",
      "\targs := m.Called(token)\n",
      "\treturn args.String(0), args.Error(1)\n",
      "}\n",
      "\n",
      "func (m *MockStorage) StoreData(token string, data string) error {\n",
      "\targs := m.Called(token, data)\n",
      "\treturn args.Error(0)\n",
      "}\n",
      "\n",
      "func TestProcessingServiceImpl_Process(t *testing.T) {\n",
      "\tmockProcessor := new(MockProcessor)\n",
      "\tmockTokenCreator := new(MockTokenCreator)\n",
      "\tmockStorage := new(MockStorage)\n",
      "\n",
      "\tservice := NewProcessingService(mockProcessor, mockTokenCreator, mockStorage)\n",
      "\n",
      "\traw := \"raw data\"\n",
      "\tprocessed := \"processed data\"\n",
      "\ttoken := \"token\"\n",
      "\n",
      "\tmockProcessor.On(\"Process\", raw).Return(processed, nil)\n",
      "\tmockTokenCreator.On(\"CreateToken\", processed).Return(token)\n",
      "\tmockStorage.On(\"StoreData\", token, processed).Return(nil)\n",
      "\n",
      "\tresult, err := service.Process(raw)\n",
      "\n",
      "\tassert.NoError(t, err)\n",
      "\tassert.Equal(t, token, result)\n",
      "\n",
      "\tmockProcessor.AssertExpectations(t)\n",
      "\tmockTokenCreator.AssertExpectations(t)\n",
      "\tmockStorage.AssertExpectations(t)\n",
      "}\n",
      "\n",
      "func TestProcessingServiceImpl_Retrieve(t *testing.T) {\n",
      "\tmockProcessor := new(MockProcessor)\n",
      "\tmockTokenCreator := new(MockTokenCreator)\n",
      "\tmockStorage := new(MockStorage)\n",
      "\n",
      "\tservice := NewProcessingService(mockProcessor, mockTokenCreator, mockStorage)\n",
      "\n",
      "\ttoken := \"token\"\n",
      "\tdata := \"data\"\n",
      "\n",
      "\tmockStorage.On(\"RetieveData\", token).Return(data, nil)\n",
      "\n",
      "\tresult, err := service.Retrieve(token)\n",
      "\n",
      "\tassert.NoError(t, err)\n",
      "\tassert.Equal(t, data, result)\n",
      "\n",
      "\tmockStorage.AssertExpectations(t)\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "print(write_unit_tests(src_dir+\"/demo/skeleton.go\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d420370",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
